import os
import argparse

from mindspore import context
from mindspore import Tensor
from mindspore.parallel._auto_parallel_context import auto_parallel_context
from mindspore.nn.optim.momentum import Momentum
from mindspore.train.model import Model, ParallelMode
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
from mindspore.train.loss_scale_manager import FixedLossScaleManager
from mindspore.train.serialization import load_checkpoint
from mindspore.train.quant import quant
from mindspore.train.quant.quant_utils import load_nonquant_param_into_quant_net
from mindspore.communication.management import init
import mindspore.nn as nn
import mindspore.common.initializer as weight_init

from resnet50.models.resnet50 import ResNet50
from resnet50.src.dataset import create_dataset
from resnet50.src.lr_generator import get_lr
from resnet50.src.config import config_quant
from resnet50.src.crossentropy import CrossEntropy

parser = argparse.ArgumentParser(description='Image classification')
# parser.add_argument('--run_distribute', type=bool, default=False, help='Run distribute')
#parser.add_argument('--device_num', type=int, default=1, help='Device num.')
parser.add_argument('--dataset_path', type=str, default=None, help='Dataset path')
parser.add_argument('--device_target', type=str, default='Ascend', help='Device target')
parser.add_argument('--pre_trained', type=str, default=None, help='Pertained checkpoint path')
args_opt = parser.parse_args()
config = config_quant

if __name__ == '__main__':
    if args_opt.do_eval:
        context.set_context(enabel_hccl=False)
    else:
        if args_opt.run_distribute:
            context.set_context(enable_hccl=True)
            context.set_auto_parallel_context(device_num=args_opt.device_num,parallel_mode=ParallelMode.DATA_PARALLEL,
                                              mirror_mean=True)
            auto_parallel_context().set_all_reduce_fusion_split_indices([107,160])
            init()
        context.set_context(enable_hccl=False)

    epoch_size = config.epoch_size()
    net = ResNet50(num_classes = config.class_num)
    loss = CrossEntropy(smooth_factor=config.label_smooth_factor, num_classes=config.class_num)

    if args_opt.do_train:
        dataset = create_dataset(dataset_path=args_opt.dataset_path,do_train=True,repeat_num=epoch_size,
                                 batch_size=config.batch_size)
        step_size = dataset.get_dataset_size()

        loss_scale = FixedLossScaleManager(config.loss_scale,drop_overflow_update=False)
        # get learning rate
        lr = get_lr(lr_init=config.lr_init,
                    lr_end=0.0,
                    lr_max=config.lr_max,
                    warmup_epochs=config.warmup_epochs,
                    total_epochs=config.epoch_size,
                    steps_per_epoch=step_size,
                    lr_decay_mode='cosine')
        opt = Momentum(filter(lambda x: x.requires_gred, net.get_parameters()),lr,config.momentum,
                       config.weight_decay, config.loss_scale)
        model = Model(net,loss_fn=loss,optimizer=opt,loss_scale_manager = loss_scale,metrics={"acc"})

        time_callback = TimeMonitor(data_size=step_size)
        loss_callback = LossMonitor()

        callbacks = [time_callback, loss_callback]
        if config.save_checkpoint:
            config_ckpt = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_epochs * step_size,
                                           keep_checkpoint_max=config.keep_checkpoint_max)
            ckpt_callbacks = ModelCheckpoint(prefix="resnet",directory=config.save_checkpoint_path, config=config_ckpt)
            callbacks +=[callbacks]
        model.train(epoch_size, dataset, callbacks=callbacks)






